%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%Plots
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
weekstr=[];
for ww=1:1:numel(Ivec)
    weekstr=[weekstr {['Week ',num2str(ww)]}];
end

%plot value functions no epi case
plot_value_functions_no_epi;

% Epi Simulation: consumption policy as function of asset grid during epi
% %plot consumption per capita by types and age thru epi
time=1:1:numel(muvec);
figure;
subplot(2,3,1)
mesh(time,grid.b(mpar.grid_sim_idx),(cmat.sy));
ylabel('b');
xlabel('weeks');
zlabel('level');
title('Consumption, sy')

subplot(2,3,4)
mesh(time,grid.b(mpar.grid_sim_idx),(cmat.so));
ylabel('b');
xlabel('weeks');
zlabel('level');
title('Consumption, so')

subplot(2,3,2)
mesh(time,grid.b(mpar.grid_sim_idx),(cmat.iy));
ylabel('b');
xlabel('weeks');
zlabel('level');
title('Consumption, iy')

subplot(2,3,5)
mesh(time,grid.b(mpar.grid_sim_idx),(cmat.io));
xlabel('b');
ylabel('time');
zlabel('level');
title('Consumption, io')

subplot(2,3,3)
mesh(time,grid.b(mpar.grid_sim_idx),(cmat.ry));
ylabel('b');
xlabel('weeks');
zlabel('level');
title('Consumption, ry')

subplot(2,3,6)
mesh(time,grid.b(mpar.grid_sim_idx),(cmat.ro));
ylabel('b');
xlabel('weeks');
zlabel('level');
title('Consumption, ro')

suptitle('Consumption by types and age in the epidemic')

orient landscape
print -dpdf -fillpage policy_functions_epi


%simulation
% %plot consumption per capita by types and age and initial assets 
% all expressed as percent deviations from no epidemic counterfactual
time=1:1:numel(muvec);
figure;
subplot(2,3,1)
plot(time,(c_EpiSim.sy./c_NoEpiSim.sy(1:numel(time))-1)*100,'b'); hold on
xlabel('weeks');
ylabel('% of no-epi baseline');
title('Consumption, sy')
axis([ 0 62 -100 5])
vline(mpar.infect_surprise_week+1)


subplot(2,3,4)
plot(time,(c_EpiSim.so./c_NoEpiSim.so(1:numel(time))-1)*100,'b'); hold on
xlabel('weeks');
ylabel('% of no-epi baseline');
title('Consumption, so')
axis([ 0 62 -100 5])
vline(mpar.infect_surprise_week+1)

subplot(2,3,2)
plot(time,(c_EpiSim.iy./c_NoEpiSim.sy(1:numel(time))-1)*100,'b'); hold on  %not that relevant no-epi type is sy
xlabel('weeks');
ylabel('% of no-epi baseline');
title('Consumption, iy')
axis([ 0 62 -100 5])
vline(mpar.infect_surprise_week+1)

subplot(2,3,5)
plot(time,(c_EpiSim.io./c_NoEpiSim.so(1:numel(time))-1)*100,'b'); hold on %not that relevant no-epi type is so
xlabel('weeks');
ylabel('% of no-epi baseline');
title('Consumption, io')
axis([ 0 62 -100 5])
vline(mpar.infect_surprise_week+1)

subplot(2,3,3)
plot(time,(c_EpiSim.ry./c_NoEpiSim.sy(1:numel(time))-1)*100,'b'); hold on %not that relevant no-epi type is sy
xlabel('weeks');
ylabel('% of no-epi baseline');
title('Consumption, ry')
axis([ 0 62 -100 5])
vline(mpar.infect_surprise_week+1)

subplot(2,3,6)
plot(time,(c_EpiSim.ro./c_NoEpiSim.so(1:numel(time))-1)*100,'b'); hold on %not that relevant no-epi type is so
xlabel('weeks');
ylabel('% of no-epi baseline');
title('Consumption, ro')
axis([ 0 62 -100 5])
vline(mpar.infect_surprise_week+1)

suptitle('Consumption by types and age in the epidemic')

orient landscape
print -dpdf -fillpage consumption_by_type_age_assets_during_epidemic




%Probs of getting infected
tauymat=par.pi1*cmat.sy.*Ivec+par.pi2.*Ivec;
tauomat=par.pi1*cmat.so.*Ivec+par.pi2.*Ivec;
    
if isempty(find(tauymat>1))==0
    error('tauy > 1');
end
if isempty(find(tauomat>1))==0
    error('tauyo > 1');
end

figure;
subplot(2,2,1)
mesh(time,grid.b(mpar.grid_sim_idx),tauymat);
ylabel('b');
xlabel('weeks');
%zlabel('% of ini. cons.');
title('tau, y')
view([-23 14]);

subplot(2,2,2)
mesh(time,grid.b(mpar.grid_sim_idx),tauomat);
ylabel('b');
xlabel('weeks');
%zlabel('% of ini. cons.');
title('tau, o')
view([-23 14]); 


%constant gain learning for cfr's
subplot(2,2,3)
plot(pidy_vec_beliefs,'r-'); hold on
plot(par.pidy_vec_true,'b--'); hold on
plot(pidy_vec_beliefs*0,'k:'); hold on
legend('Const. Gain','Steady State');
title('Weekly Prob. of Dying, Constant Gain Learning, Young');
axis tight
vline(mpar.infect_surprise_week+1)
legend('Beliefs','True','Location','Northeast');
legend boxoff

subplot(2,2,4)
plot(pido_vec_beliefs,'r-'); hold on
plot(par.pido_vec_true,'b--'); hold on
plot(pido_vec_beliefs*0,'k:'); hold on
title('Weekly Prob. of Dying, Constant Gain Learning, Old');
legend('Const. Gain','Steady State');
axis tight
vline(mpar.infect_surprise_week+1)
legend('Beliefs','True','Location','Northeast')
legend boxoff

orient landscape
print -dpdf -fillpage tauo_tauy_and_cons_gain_paths





%epi dynamics
figure;
subplot(2,4,1:2)
plot(time,Ivec,'b-');
title('Infections (Data), I');
xlabel('weeks');
vline(mpar.infect_surprise_week+1)

subplot(2,4,3:4)
plot(time,muvec,'b-');
title('Containment (Data), \mu');
xlabel('weeks');
vline(mpar.infect_surprise_week+1)

subplot(2,4,5)
plot(time,Iy);hold on;
plot(time,Io);hold on;
title('Infected (SIR implied)');
xlabel('weeks');
legend('young','old');
legend boxoff;
vline(mpar.infect_surprise_week+1)

subplot(2,4,6)
plot(time,Sy);hold on;
plot(time,So);hold on;
title('Susceptibles (SIR implied)');
xlabel('weeks');
legend('young','old');
legend boxoff;
vline(mpar.infect_surprise_week+1)

subplot(2,4,7)
plot(time,Ry);hold on;
plot(time,Ro);hold on;
title('Recovered (SIR implied)');
xlabel('weeks');
legend('young','old');
legend boxoff;
vline(mpar.infect_surprise_week+1)

subplot(2,4,8)
plot(time,Dy);hold on;
plot(time,Do);hold on;
title('Deceased (SIR implied)');
xlabel('weeks');
legend('young','old');
legend boxoff;
vline(mpar.infect_surprise_week+1)

orient landscape
print -dpdf -fillpage infections_containment_used_in_epi_simulation



%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%KEY FIGURE: EFFECTS OF EPI+CONTAINMENT ON YOUNG AND OLD
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
conflevel=0.95;
factor=-norminv((1-conflevel)/2);

figure;
sub1=subplot(1,1,1);

plot((1:1:14),tty_data.Consy_data,'b--','LineWidth',2); hold on %data mean young

grpyat = [ (1:1:14)' (tty_data.Consy_data(1:end)-factor*100*par.dat_stderr_young(1:end)'); (14:-1:1)' (tty_data.Consy_data(end:-1:1))+factor*100*par.dat_stderr_young(end:-1:1)'];
patch(grpyat(:,1),grpyat(:,2),[0.8 0.8 1],'edgecolor',[0.8 0.8 1]); hold on %data 95% young
%use uisetcolor to get RBG codes

plot((1:1:14),cons_monthly.Cy.Consy(1:end-1),'b-','LineWidth',2); hold on %model young

plot((1:1:14),tto_data.Conso_data,'r--','LineWidth',2);%data mean old

grpyat = [ (1:1:14)' (tto_data.Conso_data(1:end)-factor*100*par.dat_stderr_old(1:end)'); (14:-1:1)' (tto_data.Conso_data(end:-1:1))+factor*100*par.dat_stderr_old(end:-1:1)'];
patch(grpyat(:,1),grpyat(:,2),[1 0.8 0.8],'edgecolor',[1 0.8 0.8]); hold on%data 95% old

plot((1:1:14),cons_monthly.Co.Conso(1:end-1),'r-','LineWidth',2); hold on%model old

plot((1:1:14),tty_data.Consy_data,'b--','LineWidth',2); hold on %data mean young
plot((1:1:14),cons_monthly.Cy.Consy(1:end-1),'b-','LineWidth',2); hold on %model young
plot((1:1:14),tto_data.Conso_data,'r--','LineWidth',2);%data mean old

plot((1:1:14),0*tto_data.Conso_data,'k:','LineWidth',1.5);
titl=title('Consumption of Young and Old in the Pandemic','FontSize',12);
legend1=legend('Data: Young (Mean)','Data: Young (95%)', 'Model: Young', 'Data: Old (Mean)','Data: Old (95%)', 'Model: Old','Location','Southwest','FontSize',12);
box off;
legend box off;
ylabel('% deviations from no epidemic baseline','FontSize',12);
set(legend1,...
    'Position',[0.317548452641451 0.250788781770377 0.264285714285714 0.229761904761905],...
    'Orientation','vertical');
sub1.XLim(1)=1;%'01-Mar-2020';
sub1.XLim(2)=14;%'01-Mar-2021';
sub1.YLim(1)=-50;
sub1.YLim(2)=15;

datstring=['Mar 2020'
    'Apr 2020'
    'May 2020'
    'Jun 2020'
    'Jul 2020'
    'Aug 2020'
    'Sep 2020'
    'Oct 2020'
    'Nov 2020'
    'Dec 2020'
    'Jan 2021'
    'Feb 2021'
    'Mar 2021'
    'Apr 2021'];

%xticks(cons_monthly.Cy.Time(1:end-1));
xticks(1:14);
%xticklabels({cons_monthly.Cy.Time(1:end-1)});
xticklabels(datstring)
xtickangle(-35);
set(gca,'FontSize',12);
set(titl,'FontSize',18);

orient landscape
print -dpdf -fillpage cons_young_old_epidemic


%print out results
disp(' ');

disp('Model vs. Data Consumption (in percent of no-epi baseline)');

%young, data 
%first wave
young_first_data=min(tty_data.Consy_data(1:7))
%second wage
young_second_data=min(tty_data.Consy_data(8:end))

%young, model
%first wave
young_first_model=min(cons_monthly.Cy.Consy(1:7))
%second wave
young_second_model=min(cons_monthly.Cy.Consy(8:end))

%old, data 
%first wave
old_first_data=min(tto_data.Conso_data(1:7))
%second wage
old_second_data=min(tto_data.Conso_data(8:end))


%old, model
%first wave
old_first_model=min(cons_monthly.Co.Conso(1:7))
%second wave
old_second_model=min(cons_monthly.Co.Conso(8:end))
